import tensorflow as tf
import os
from nets.nets_factory import get_network_fn, arg_scopes_map
import cv2
import numpy as  np
from PIL import Image
from matplotlib import pyplot as plt
from preprocessing import preprocessing_factory


os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
slim = tf.contrib.slim


labels_file = "./labels.txt"


def preprocess_image(image):
    image = tf.io.read_file(image)
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.convert_image_dtype(image, dtype=tf.float32)  # 如果是float 进来就不处理了
    image = tf.image.central_crop(image, central_fraction=0.875)  # 此操作后size就变了
    # # Resize the image to the original height and width.
    image = tf.expand_dims(image, 0)
    image = tf.image.resize_bilinear(image, [299, 299], align_corners=False)
    # image = tf.squeeze(image, [0])
    # Finally, rescale to [-1,1] instead of [0, 1)
    image = tf.subtract(image, 0.5)
    image = tf.multiply(image, 2.0)
    with tf.Session() as sess:

        return sess.run(image)


def processing_with_pil(img_file,show=True,central_fraction=0.875,resize=(299,299)):
    img = Image.open(img_file)
    if resize:
        img = img.resize(resize)
    shape = img.size

    w,h = shape[0], shape[1]
    wd = w - int((w - w * central_fraction) / 2) * 2
    hd = h - int((h - h * central_fraction) / 2) * 2
    img = img.crop(((w-wd)/2,(h-hd)/2,(w-wd)/2+ wd, (h-hd)/ 2 + hd))  # left, up, right, below
    img = img.convert('RGB')
    if resize:
        img = img.resize(resize)
    img = np.array(img)

    img_normal = (img / 255 - 0.5) / 2
    # img_normal = img_normal.astype(np.int8)
    print(img_normal.shape)
    if show:
        plt.imshow(img_normal)
        plt.show()
    return img_normal.reshape([1,299,299,3])


def processing_with_cv(img_file,central_fraction=0.875,show=False):
    img = cv2.imread(img_file)
    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)

    ## crop the image
    w,h = img.shape[0],img.shape[1]
    wd = int(w * central_fraction)
    hd = int(h * central_fraction)
    img = img[
        int((w - wd) / 2):int((w - wd) / 2 + wd),
        int((h - hd) / 2):int((h - hd) / 2 + hd),:]

    img = cv2.resize(img, (299, 299))

    img = (img / 255 - 0.5) / 2

    if show:
        cv2.imshow("",img)
        cv2.waitKey(0)

    return img.reshape([1,299,299,3])#.astype(np.int8)


img_file = "./test_video/t/"

img_l = os.listdir(img_file)[:]


def main():
    # with tf.Graph().as_default() as g:
    f = open(labels_file,'r')
    lines = f.readlines()
    names = {}  # {0: 'normal', 1: 'porn', 2: 'sexy'}
    for line in lines:
        i, label = line.split(':')
        names[int(i)] = label.strip()

    names_tensor = tf.constant(list(names.values()))
    names_lookup_table = tf.contrib.lookup.index_to_string_table_from_tensor(names_tensor)

    jpegs = tf.placeholder(tf.float32,shape=[None,299,299,3],name='input')
    # img = tf.map_fn(preprocess_image, jpegs, dtype=tf.float32)


    net_fn = get_network_fn('inception_v4',3,is_training=False)
    with slim.arg_scope(arg_scopes_map['inception_v4']()):  # resnet_arg_scope
        logits, end_points = net_fn(jpegs)
        probs = tf.nn.softmax(logits)

        topk_probs, topk_indices = tf.nn.top_k(probs, 1)

        # todo：不知道什么用
        topk_names = names_lookup_table.lookup(tf.to_int64(topk_indices))


    with tf.compat.v1.Session() as sess:

        saver = tf.train.Saver()
        ckpt = tf.train.get_checkpoint_state("./my-data/train_sj_d3_3/")
        # print(ckpt.model_checkpoint_path)
        saver.restore(sess,os.path.join("./my-data/train_sj_001/","model.ckpt-234972"))

        # input_node_name = sess.graph.get_tensor_by_name('input:0')
        # output_node_name = sess.graph.get_tensor_by_name('InceptionV4/Logits/Predictions:0')
        # res = sess.run(output_node_name, feed_dict={input_node_name: [tf.gfile.FastGFile(img_file, 'rb').read()]}) #发送二进制是成功啦 [[4.9104528e-06 9.9984705e-01 1.4797678e-04]]
        f = open("res/t.txt",'w')

        for img in img_l:
        # for img in range(1,len(img_l)):

            if os.path.isfile(img):
                continue
            img_fil = img_file + img
            print(img)
            res = sess.run([probs,topk_indices], feed_dict={jpegs:preprocess_image(img_fil)})
            log = str(img) + "  " + names[int(res[1])] + ": " + str(res[0])+"\n"

            print(log)
            f.write(log)
            f.flush()
        f.close()


if __name__ == '__main__':
    main()
    # img = img_file + img_l[0]
    #
    # image_preprocessing_fn = preprocessing_factory.get_preprocessing("resnet_v2_50",is_training=False)
    # image = image_preprocessing_fn(cv2.imread(img),224,224)
    # print(image)
    # with tf.Session() as sess:
    #     img = sess.run(img)
    #     print(img)

